#include "StdAfx.hpp"
#include "KDTree.hpp"
#include "Scene.hpp"

using namespace common;

namespace {

const uint NODE_POOL_BLOCK_CAPACITY = 512;
const uint OBJ_LINK_POOL_BLOCK_CAPACITY = 1024;

const uint NODE_OBJ_COUNT_UPPER_LIMIT = 32;
const uint NODE_OBJ_COUNT_LOWER_LIMIT = 16;
const float LOOSE_FACTOR = 0.1f;

} // namespace anonymous


KDTree::KDTree(common::BOX &WorldBox)
: m_NodePool(NODE_POOL_BLOCK_CAPACITY)
, m_ObjLinkPool(OBJ_LINK_POOL_BLOCK_CAPACITY)
{
	ZeroMemory(m_NodesPerLevel, sizeof(m_NodesPerLevel));
	ZeroMemory(m_ObjectsPerLevel, sizeof(m_ObjectsPerLevel));

	m_RootNode = AllocNode();
	
	m_RootNode->Level = 0;
	m_RootNode->Parent = NULL;
	m_RootNode->Children[0] = NULL;
	m_RootNode->Children[1] = NULL;
	m_RootNode->Box = WorldBox;
	m_RootNode->FirstObjLink = NULL;

	m_NodesPerLevel[0] = 1;
}

KDTree::~KDTree()
{
	// Pools will automatically free their elements.
}

void KDTree::InsertObject(SceneObject *obj)
{
	assert(obj->m_KDNode == NULL);

	const BOX &objWorldBox = obj->GetWorldBoundingBox();
	KDNode *node = FindNodeForObject(objWorldBox);
	AddObjectToNode(node, obj);
	if (IsNodeLeaf(node) &&
		node->Level < KD_TREE_MAX_DEPTH-1 &&
		GetNodeObjectCount(node) > NODE_OBJ_COUNT_UPPER_LIMIT)
	{
		SplitNode(node);
	}
}

void KDTree::RemoveObject(SceneObject *obj)
{
	KDNode *node = obj->m_KDNode;
	assert(node != NULL);

	RemoveObjectFromNode(node, obj);
	if (node->Parent != NULL)
		TryMergeNode(node->Parent);

	obj->m_KDNode = NULL;
}

void KDTree::ObjectTransformChange(SceneObject *obj)
{
	KDNode *node = obj->m_KDNode;
	assert(node != NULL);
	const BOX &objWorldBox = obj->GetWorldBoundingBox();
	if (BoxInBox(objWorldBox, node->Box))
		return;
	RemoveObject(obj);
	InsertObject(obj);
}

void KDTree::ListObjectsIntersectingFrustum(SceneObjectPtrVector &out, const common::FRUSTUM_PLANES &Frustum, SceneObjectPredicate *Pred)
{
	out.clear();
	AddObjectsIntersectingFrustum(out, m_RootNode, Frustum, Pred);
}

void KDTree::ListObjectsIntersectingSweptBox(SceneObjectPtrVector &out, const common::BOX &box, const common::VEC3 &boxDir, SceneObjectPredicate *Pred)
{
	out.clear();
	AddObjectsIntersectingSweptBox(out, m_RootNode, box, boxDir, Pred);
}

/*void KDTree::ListObjectsIntersectingBox(SceneObjectPtrVector &out, const common::BOX &Box, SceneObjectPredicate *Pred)
{
	out.clear();
	AddObjectsIntersectingBox(out, m_RootNode, Box, Pred);
}*/

void KDTree::ListObjectsIntersectingSphere(SceneObjectPtrVector &out, const common::VEC3 &sphereCenter, float sphereRadius, SceneObjectPredicate *pred)
{
	out.clear();
	AddObjectsIntersectingSphere(out, m_RootNode, sphereCenter, sphereRadius, pred);
}

void KDTree::AddObjectToNode(KDNode *node, SceneObject *obj)
{
	obj->m_KDNode = node;

	KDObjLink *newObjLink = AllocObjLink();
	newObjLink->Obj = obj;
	newObjLink->NextObjLink = node->FirstObjLink;
	node->FirstObjLink = newObjLink;

	m_ObjectsPerLevel[node->Level]++;
}

void KDTree::RemoveObjectFromNode(KDNode *node, SceneObject *obj)
{
	assert(obj->m_KDNode == node);

	KDObjLink *objLink = node->FirstObjLink, *prevLink = NULL;
	while (objLink)
	{
		if (objLink->Obj == obj)
		{
			DeleteObjLink(node, prevLink, objLink);
			obj->m_KDNode = NULL;
			return;
		}
		prevLink = objLink;
		objLink = objLink->NextObjLink;
	}

	assert(0 && "Object not found in KDNode");
}

void KDTree::CalcOptimalSplit(uint *outNormalAxisIndex, float *outSplitV, KDNode *node)
{
	// Simple algorithm
	/*{
		*outNormalAxisIndex = node->Level % 3;
		*outSplitV = (node->Box.Min[*outNormalAxisIndex] + node->Box.Max[*outNormalAxisIndex]) * 0.5f;
	}*/

	// Complex algorithm
	{
		VEC3 Mean, Variance;
		//CalcNodeObjectBoundingBoxesMeanAndVariance(&Mean, &Variance, node);
		CalcNodeObjectPositionsMeanAndVariance(&Mean, &Variance, node);
		
		if (Variance.x > Variance.y && Variance.x > Variance.z)
			*outNormalAxisIndex = 0; // Cut with YZ plane
		else if (Variance.y > Variance.z)
			*outNormalAxisIndex = 1; // Cut with XZ plane
		else
			*outNormalAxisIndex = 2; // Cut with XY plane

		float splitVFromMean = Mean[*outNormalAxisIndex];
		float splitVFromBox = (node->Box.Min[*outNormalAxisIndex] + node->Box.Max[*outNormalAxisIndex]) * 0.5f;
		*outSplitV = (splitVFromMean + splitVFromBox) * 0.5f;
	}
}

void KDTree::SplitNode(KDNode *node)
{
	assert(IsNodeLeaf(node));

	uint NormalAxisIndex;
	float SplitV;
	CalcOptimalSplit(&NormalAxisIndex, &SplitV, node);

	uint newLevel = node->Level + 1;
	m_NodesPerLevel[newLevel] += 2;

	KDNode *n0 = AllocNode();
	KDNode *n1 = AllocNode();
	node->Children[0] = n0;
	node->Children[1] = n1;
	n0->Level = n1->Level = newLevel;
	n0->Parent = n1->Parent = node;
	n0->FirstObjLink = n1->FirstObjLink = NULL;
	n0->Children[0] = n0->Children[1] = NULL;
	n1->Children[0] = n1->Children[1] = NULL;

	n0->Box = n1->Box = node->Box;
	n0->Box.Max[NormalAxisIndex] = n1->Box.Min[NormalAxisIndex] = SplitV;

	float parentBoxSize = node->Box.Max[NormalAxisIndex] - node->Box.Min[NormalAxisIndex];
	float extendDist = parentBoxSize * LOOSE_FACTOR;

	n0->Box.Max[NormalAxisIndex] += extendDist;
	n1->Box.Min[NormalAxisIndex] -= extendDist;

	TransferObjectsToChildNodes(node);
}

void KDTree::MergeNode(KDNode *node)
{
	assert(!IsNodeLeaf(node));
	assert(IsNodeLeaf(node->Children[0]) && IsNodeLeaf(node->Children[1]));

	KDObjLink *prevLink, *currLink, *nextLink;
	
	prevLink = NULL;
	currLink = node->Children[0]->FirstObjLink;
	while (currLink)
	{
		nextLink = currLink->NextObjLink;

		AddObjectToNode(node, currLink->Obj);
		DeleteObjLink(node->Children[0], prevLink, currLink);

		//prevLink = currLink;
		currLink = nextLink;
	}

	prevLink = NULL;
	currLink = node->Children[1]->FirstObjLink;
	while (currLink)
	{
		nextLink = currLink->NextObjLink;

		AddObjectToNode(node, currLink->Obj);
		DeleteObjLink(node->Children[1], prevLink, currLink);

		//prevLink = currLink;
		currLink = nextLink;
	}

	m_NodesPerLevel[node->Children[0]->Level] -= 2;

	FreeNode(node->Children[0]); node->Children[0] = NULL;
	FreeNode(node->Children[1]); node->Children[1] = NULL;
}

uint KDTree::GetNodeObjectCount(KDNode *node)
{
	uint objCount = 0;
	KDObjLink *objLink = node->FirstObjLink;
	while (objLink)
	{
		objCount++;
		objLink = objLink->NextObjLink;
	}
	return objCount;
}

KDNode * KDTree::FindNodeForObject(const common::BOX &objWorldBox)
{
	KDNode *node = m_RootNode;
	while (!IsNodeLeaf(node))
	{
		if (BoxInBox(objWorldBox, node->Children[0]->Box))
			node = node->Children[0];
		else if (BoxInBox(objWorldBox, node->Children[1]->Box))
			node = node->Children[1];
		else
			break;
	}
	return node;
}

void KDTree::CalcNodeObjectBoundingBoxesMeanAndVariance(common::VEC3 *outMean, common::VEC3 *outVariance, KDNode *node)
{
	float objCount = 0.f;
	KDObjLink *objLink = node->FirstObjLink;
	*outMean = VEC3_ZERO;
	while (objLink)
	{
		const BOX &objWorldBox = objLink->Obj->GetWorldBoundingBox();
		*outMean += objWorldBox.Min + objWorldBox.Max;
		objCount += 1.f;
		objLink = objLink->NextObjLink;
	}
	float objCount_2_rcp = 0.5f / objCount;
	*outMean *= objCount_2_rcp;

	*outVariance = VEC3_ZERO;
	objLink = node->FirstObjLink;
	VEC3 tmp1, tmp2;
	while (objLink)
	{
		const BOX &objWorldBox = objLink->Obj->GetWorldBoundingBox();
		
		tmp1 = objWorldBox.Min - *outMean;
		tmp2 = objWorldBox.Max - *outMean;
		outVariance->x += tmp1.x * tmp1.x + tmp2.x * tmp2.x;
		outVariance->y += tmp1.y * tmp1.y + tmp2.y * tmp2.y;
		outVariance->z += tmp1.z * tmp1.z + tmp2.z * tmp2.z;
		
		objLink = objLink->NextObjLink;
	}
	*outVariance *= objCount_2_rcp;
}

void KDTree::CalcNodeObjectPositionsMeanAndVariance(common::VEC3 *outMean, common::VEC3 *outVariance, KDNode *node)
{
	float objCount = 0.f;
	KDObjLink *objLink = node->FirstObjLink;
	*outMean = VEC3_ZERO;
	while (objLink)
	{
		*outMean += objLink->Obj->GetWorldPosition();
		objCount += 1.f;
		objLink = objLink->NextObjLink;
	}
	float objCountRcp = 1.f / objCount;
	*outMean *= objCountRcp;

	*outVariance = VEC3_ZERO;
	objLink = node->FirstObjLink;
	VEC3 tmp;
	while (objLink)
	{
		const BOX &objWorldBox = objLink->Obj->GetWorldBoundingBox();
		
		tmp = objLink->Obj->GetWorldPosition() - *outMean;
		outVariance->x += tmp.x * tmp.x;
		outVariance->y += tmp.y * tmp.y;
		outVariance->z += tmp.z * tmp.z;
		
		objLink = objLink->NextObjLink;
	}
	*outVariance *= objCountRcp;
}

void KDTree::TransferObjectsToChildNodes(KDNode *node)
{
	assert(!IsNodeLeaf(node));

	KDObjLink *objLink = node->FirstObjLink, *prevLink = NULL, *nextLink;
	while (objLink)
	{
		nextLink = objLink->NextObjLink;

		SceneObject *obj = objLink->Obj;
		const BOX &objWorldBox = obj->GetWorldBoundingBox();
		if (BoxInBox(objWorldBox, node->Children[0]->Box))
		{
			DeleteObjLink(node, prevLink, objLink);
			AddObjectToNode(node->Children[0], obj);
		}
		else if (BoxInBox(objWorldBox, node->Children[1]->Box))
		{
			DeleteObjLink(node, prevLink, objLink);
			AddObjectToNode(node->Children[1], obj);
		}
		else
			prevLink = objLink;

		objLink = nextLink;
	}
}

void KDTree::DeleteObjLink(KDNode *node, KDObjLink *prevLink, KDObjLink *currLink)
{
	m_ObjectsPerLevel[node->Level]--;

	if (prevLink != NULL)
		prevLink->NextObjLink = currLink->NextObjLink;
	else
		node->FirstObjLink = currLink->NextObjLink;
	
	FreeObjLink(currLink);
}

void KDTree::ValidateNode(KDNode *node)
{
	assert(node->Box.IsValid());
	assert(node->Parent || node == m_RootNode);
	assert( (node->Children[0] == 0) == (node->Children[1] == 0) );

	KDObjLink *link = node->FirstObjLink;
	while (link)
	{
		assert(link->Obj->m_KDNode == node);
		link = link->NextObjLink;
	}

	if (node->Children[0])
	{
		ValidateNode(node->Children[0]);
		ValidateNode(node->Children[1]);
	}
}

void KDTree::AddObjectsIntersectingFrustum(SceneObjectPtrVector &out, KDNode *node, const common::FRUSTUM_PLANES &Frustum, SceneObjectPredicate *Pred)
{
	KDObjLink *link = node->FirstObjLink;
	while (link)
	{
		if (Pred == NULL || Pred->Test(link->Obj))
			if (BoxToFrustum_Fast(link->Obj->GetWorldBoundingBox(), Frustum))
				out.push_back(link->Obj);
		link = link->NextObjLink;
	}

	if (!IsNodeLeaf(node))
	{
		if (BoxToFrustum_Fast(node->Children[0]->Box, Frustum))
			AddObjectsIntersectingFrustum(out, node->Children[0], Frustum, Pred);
		if (BoxToFrustum_Fast(node->Children[1]->Box, Frustum))
			AddObjectsIntersectingFrustum(out, node->Children[1], Frustum, Pred);
	}
}

void KDTree::AddObjectsIntersectingSweptBox(SceneObjectPtrVector &out, KDNode *node, const common::BOX &box, const common::VEC3 &boxDir, SceneObjectPredicate *Pred)
{
	float t1, t2;
	for (KDObjLink *link = node->FirstObjLink; link; link = link->NextObjLink)
		if (Pred == NULL || Pred->Test(link->Obj))
			if (SweptBoxToBox(box, link->Obj->GetWorldBoundingBox(), boxDir, &t1, &t2) && t1 >= 0.f)
				out.push_back(link->Obj);

	if (!IsNodeLeaf(node))
	{
		if (SweptBoxToBox(box, node->Children[0]->Box, boxDir, &t1, &t2) && t1 >= 0.f)
			AddObjectsIntersectingSweptBox(out, node->Children[0], box, boxDir, Pred);
		if (SweptBoxToBox(box, node->Children[1]->Box, boxDir, &t1, &t2) && t1 >= 0.f)
			AddObjectsIntersectingSweptBox(out, node->Children[1], box, boxDir, Pred);
	}
}

void KDTree::AddObjectsIntersectingSphere(SceneObjectPtrVector &out, KDNode *node, const common::VEC3 &sphereCenter, float sphereRadius, SceneObjectPredicate *pred)
{
	for (KDObjLink *link = node->FirstObjLink; link; link = link->NextObjLink)
		if (pred == NULL || pred->Test(link->Obj))
			if (SphereToBox(sphereCenter, sphereRadius, link->Obj->GetWorldBoundingBox()))
				out.push_back(link->Obj);

	if (!IsNodeLeaf(node))
	{
		if (SphereToBox(sphereCenter, sphereRadius, node->Children[0]->Box))
			AddObjectsIntersectingSphere(out, node->Children[0], sphereCenter, sphereRadius, pred);
		if (SphereToBox(sphereCenter, sphereRadius, node->Children[1]->Box))
			AddObjectsIntersectingSphere(out, node->Children[1], sphereCenter, sphereRadius, pred);
	}
}

void KDTree::TryMergeNode( KDNode *node )
{
	if (!IsNodeLeaf(node) &&
		IsNodeLeaf(node->Children[0]) &&
		IsNodeLeaf(node->Children[1]) &&
		(GetNodeObjectCount(node) + GetNodeObjectCount(node->Children[0]) + GetNodeObjectCount(node->Children[1]) < NODE_OBJ_COUNT_LOWER_LIMIT))
	{
		MergeNode(node);
	}
}